import pinocchio as pin
from pinocchio.utils import rand, zero
# import meshcat

# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src"))
sys.path.append(dirname(dirname(__file__)))
from src import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 获取URDF模型路径
URDF_FILE_PATH = filePath(557)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
print('URDF_FILE_PATH is ', URDF_FILE_PATH)
# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)

robot.model.gravity.linear[:2] = 0.0
robot.model.gravity.linear[2] = -0.0
q = pin.neutral(robot.model)
vq = rand(robot.model.nv)
aq = zero(robot.model.nv)
t_q = zero(robot.model.nv)
pin.framesForwardKinematics(robot.model, robot.data, q)
pin.updateFramePlacements(robot.model, robot.data) 
robot.initViewer(loadModel=True)
viz.display(q)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

print('q is ', q)
print('q shape is ', q.shape)

pelvis_idx = robot.index("pelvis")
print(robot.model.names[1])

for idx, joint_name in enumerate(robot.model.names):
    print(f"Index: {idx}, Joint Name: {joint_name}")


# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
## torso
q[1] += 0.1
q[2] += -0.1
q[4] += 0


## left shoulder arm
q[5] += -0.1
q[6] += -0.8
q[8] += -0.3
q[9] += -0
q[10] += -1.3
q[11] += 0
q[12] += 0


## right shoulder arm
q[13] += 0.1
q[14] += 0.8
q[15] += 0.5
q[16] += 0.3
q[17] += 0
q[18] += 2
q[19] += 0
q[20] += 0

## head
q[21] += 0
q[22] += 0
q[23] += -0.1
q[24] += 0
q[25] += 0.1


## left leg
q[26] += 0
q[27] += -0.5
q[28] += 0
q[29] += 2
q[30] += 0.1
q[31] += 0
q[32] += -0.2
q[33] += 0
q[34] += 0


## right leg
q[35] += 0
q[36] += -0.5
q[37] += -0
q[38] += 2
q[39] += -0.1
q[40] += 0
q[41] += -0.2
q[42] += 0
q[43] += 0



viz.display(q)
# print('pelvis_idx is ', pelvis_idx)
# pelvis_placement = robot.data.oMi[pelvis_idx].copy()
# # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
# # 前向运动学
# init_config = Init_Config(ROBOT=True)

# idx = robot.index(init_config.joint_names[1])

# q1 = np.array([0, 0, 0, -1, 0, 0])

# # 这里是给定各个关节角和特定关节index，然后求这个特定关节的姿态
# # 该方法会有返回，可以写成placement=robot.placement(q0, idx)，但是这样再直接修改placement就会改变robot姿态
# # 使用Placement同时就会改变robot姿态
# robot.placement(q1, idx)
# placement = robot.data.oMi[idx].copy()
# mprint(placement)
# print('type of placement is ', type(placement))
# # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
